from gremlin_python import statics
from gremlin_python.process.anonymous_traversal import traversal
from gremlin_python.process.graph_traversal import __
from gremlin_python.process.strategies import *
from gremlin_python.driver.driver_remote_connection import DriverRemoteConnection
from gremlin_python.structure.graph import Graph
from gremlin_python.process.traversal import T
from gremlin_python.process.traversal import Order
from gremlin_python.process.traversal import Cardinality
from gremlin_python.process.traversal import Column
from gremlin_python.process.traversal import Direction
from gremlin_python.process.traversal import Operator
from gremlin_python.process.traversal import P
from gremlin_python.process.traversal import TextP
from gremlin_python.process.traversal import Pop
from gremlin_python.process.traversal import Scope
from gremlin_python.process.traversal import Barrier
from gremlin_python.process.traversal import Bindings
from gremlin_python.process.traversal import WithOptions

g = Graph().traversal().withRemote(DriverRemoteConnection('ws://211.71.15.48:8182/gremlin', 'g'))

# 以下代码为找到特定节点向下追溯的代码（更新：已经修改为向上追溯）
"""
加一个判断，起始是for节点后用另一个方法（更新：目前的方法中无需）
"""
a = g.V().has('code', TextP.containing('torch.no_grad()')).toList()
simplePaths = g.V().has('code', TextP.containing('torch.no_grad()')).repeat(__.in_().simplePath()).until(
    __.not_(__.in_('cfg'))).path().toList()
### 改进：可以添加或条件呢：“。。。或遇到所需的node就结束”

##上面这行代码中out换in_就可以向上追溯
### bug to fix ：结尾为循环时无法生成路径 （已修）

"""
GPT解决方案（弃用）
您可以使用notor步骤来将两个条件组合起来，并在循环中检查节点是否已经访问过。以下是修改后的代码示例：
g.V().has('code', TextP.containing('return module_output')).has('ast', TextP.containing('Return')).repeat(__.in_().simplePath()).until(__.or_(__.not_(__.in_('cfg')), __.not_(__.has('visited')))).sideEffect(__.has('visited', True)).path().toList()

在上面的代码中，我们使用sideEffect步骤将visited属性设置为True，以标记节点已经访问过。
请注意，这段代码假设您的图数据库中的节点具有名为visited的布尔属性来区分已访问节点和未访问节点。如果您的图数据库的节点属性和结构有所不同，请相应地修改代码。
"""

completePaths = []

##防止整个代码都是圈

if len(simplePaths) == 0 and len(a) != 0:
    simplePaths = g.V().has('code', TextP.containing('running_loss = 0.0')).has('ast', TextP.containing(
        "Assign(targets=[Name(id='running_loss'")).repeat(__.in_('cfg')).until(__.cyclicPath()).path().toList()

for path in simplePaths:
    need_merge = False
    for node in path:
        node_id = node.id
        node_ast = g.V(node_id).ast.toList()
        if "For" in node_ast[0]:
            need_merge = True
            cyclePaths = g.V(node_id).repeat(__.in_('cfg')).until(__.cyclicPath()).path().toList()
            ## 上面这行代码中out换in就可以向上追溯
            for cyclePath in cyclePaths:
                completePath = list(path) + list(cyclePath)
                completePaths.append(completePath)
    if not need_merge:
        completePaths.append(list(path))


## 去重
def remove_duplicate_paths(paths):
    return [list(x) for x in set(tuple(x) for x in paths)]


paths = remove_duplicate_paths(completePaths)

for path in paths:
    is_right = False
    print(path)
    for node in path:
        node_id = node.id
        node_ast = g.V(node_id).ast.toList()[0]
        node_code = g.V(node_id).code.toList()[0]
        node_lineo = g.V(node_id).lineo.toList()[0]
        if 'eval' in node_code:  ## 判断是否有符合要求的节点
            is_right = not is_right
    if not is_right:
        print('代码有问题')



